import os

os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
os.environ['CUDA_VISIBLE_DEVICES'] = '2'
import numpy as np
from torch import nn
from torchvision.transforms import InterpolationMode
from torch.utils.data import DataLoader
import json
import random
import torch
from torchvision import transforms
from utils import test
from dataset import CustomImageDataset
from model import swin_base_patch4_window12_384_in22k as Model
from config_hyperparam import cfg
from sklearn.manifold import TSNE
import seaborn as sns
import tqdm
import sys
import matplotlib.pyplot as plt
from sklearn.metrics import classification_report
    
@torch.no_grad()
def get_tsne_data(model, data_loader, device):


    feats = []
    labels = []

    model.eval()

    for step, data in enumerate(data_loader):
        images, labels1 = data
        images, labels1 = images.to(device), labels1.to(device)
        feat = model(images)
        feats.append(torch.flatten(feat).cpu().numpy())
        labels.append(labels1.squeeze().cpu().numpy())
    np.save('feats.npy', feats)
    np.save('labels.npy', labels)



def main():
    deviceIds = cfg.gpu_idx

    device = torch.device(cfg.device)
    nw = min([os.cpu_count(), cfg.batch_size if cfg.batch_size > 1 else 0, 8])  # number of workers

    model = Model(cfg.num_class).to(cfg.device)
    #model = nn.DataParallel(model, device_ids=deviceIds)

    test_transform = transforms.Compose([transforms.Resize((256, 256), InterpolationMode.BILINEAR),
                                         transforms.CenterCrop((cfg.resize, cfg.resize)),
                                         transforms.ToTensor(),
                                         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])

    test_dataset = CustomImageDataset(root_dir=cfg.dataset_path, transform=test_transform, train=False)

    test_loader = DataLoader(test_dataset,
                                 batch_size=1,
                                 shuffle=True,
                                 pin_memory=True,
                                 drop_last=True,
                                 num_workers=nw)


    predictions, true_labels = test(model=model, data_loader=test_loader, device=device)
    
    print(classification_report(true_labels, predictions, digits=3))


if __name__ == '__main__':
    main()
